package com.chatmcp.mcprouter.service.impl;

import com.chatmcp.mcprouter.dto.ProxyInfoDto;
import com.chatmcp.mcprouter.dto.ServerConfigDto;
import com.chatmcp.mcprouter.service.McpClientService;
import com.chatmcp.mcprouter.service.McpServerService;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.Map;
import java.util.concurrent.*;
import java.util.function.Consumer;

/**
 * 进程管理服务实现类
 * 专注于本地进程管理和进程共享机制的改进
 */
@Slf4j
@Service
public class ProcessManagementServiceImpl {

    private final ObjectMapper objectMapper;
    private final McpServerService mcpServerService;
    
    // 客户端缓存，键为serverKey，值为客户端实例
    private final Map<String, McpClientService> clientCache = new ConcurrentHashMap<>();
    
    // 进程健康检查任务映射，键为serverKey，值为调度任务Future
    private final Map<String, ScheduledFuture<?>> healthCheckTasks = new ConcurrentHashMap<>();
    
    // 任务执行器
    private final ScheduledExecutorService taskExecutor = Executors.newScheduledThreadPool(2);

    @Autowired
    public ProcessManagementServiceImpl(ObjectMapper objectMapper, McpServerService mcpServerService) {
        this.objectMapper = objectMapper;
        this.mcpServerService = mcpServerService;
    }
    
    /**
     * 获取或创建客户端实例
     * 支持进程共享和健康检查
     */
    public McpClientService getOrCreateClient(String serverKey, ServerConfigDto serverConfig) {
        McpClientService client = clientCache.get(serverKey);
        
        // 检查客户端是否存在且仍然连接
        if (client != null && client.isConnected()) {
            log.debug("Using existing client for server: {}", serverKey);
            return client;
        }
        
        // 如果客户端不存在或未连接，则创建新客户端
        try {
            client = new McpClientServiceImpl(objectMapper, null);
            boolean connected = client.connect(serverConfig);
            
            if (connected) {
                // 存储到缓存
                clientCache.put(serverKey, client);
                log.debug("Created new client for server: {}", serverKey);
                
                // 设置进程健康监视器
                if (serverConfig.getCommand() != null && !serverConfig.getCommand().isEmpty()) {
                    scheduleProcessHealthCheck(serverKey, client);
                }
                
                return client;
            } else {
                log.warn("Failed to connect to MCP server: {}", client.getError());
                return null;
            }
        } catch (Exception e) {
            log.error("Error creating MCP client: {}", e.getMessage());
            return null;
        }
    }
    
    /**
     * 调度进程健康检查任务
     * 定期检查进程健康状态，如果进程已终止则清理缓存
     */
    private void scheduleProcessHealthCheck(String serverKey, McpClientService client) {
        // 先取消已有的健康检查任务
        cancelProcessHealthCheck(serverKey);
        
        // 创建新的健康检查任务
        ScheduledFuture<?> future = taskExecutor.scheduleAtFixedRate(() -> {
            try {
                if (!client.isConnected()) {
                    log.warn("Process for server {} is no longer connected, removing from cache", serverKey);
                    clientCache.remove(serverKey);
                    cancelProcessHealthCheck(serverKey);
                    
                    // 这里也可以添加重连逻辑
                    // tryReconnectClient(serverKey);
                }
            } catch (Exception e) {
                log.error("Error checking process health: {}", e.getMessage());
            }
        }, 60, 60, TimeUnit.SECONDS);
        
        // 保存任务引用
        healthCheckTasks.put(serverKey, future);
        log.debug("Scheduled process health check for server: {}", serverKey);
    }
    
    /**
     * 取消进程健康检查任务
     */
    private void cancelProcessHealthCheck(String serverKey) {
        ScheduledFuture<?> future = healthCheckTasks.remove(serverKey);
        if (future != null && !future.isDone()) {
            future.cancel(false);
            log.debug("Cancelled process health check for server: {}", serverKey);
        }
    }
    
    /**
     * 设置通知回调
     */
    public void setupNotificationCallback(McpClientService client, Consumer<String> callback) {
        if (client != null) {
            client.setNotificationCallback(callback);
        }
    }
    
    /**
     * 关闭所有客户端连接
     * 用于应用程序关闭时清理资源
     */
    public void shutdown() {
        log.info("Shutting down process management service");
        
        // 取消所有健康检查任务
        for (String serverKey : healthCheckTasks.keySet()) {
            cancelProcessHealthCheck(serverKey);
        }
        
        // 关闭所有客户端连接
        for (Map.Entry<String, McpClientService> entry : clientCache.entrySet()) {
            try {
                entry.getValue().close();
                log.debug("Closed client for server: {}", entry.getKey());
            } catch (Exception e) {
                log.error("Error closing client for server {}: {}", entry.getKey(), e.getMessage());
            }
        }
        
        // 清空客户端缓存
        clientCache.clear();
        
        // 关闭任务执行器
        taskExecutor.shutdown();
        try {
            if (!taskExecutor.awaitTermination(5, TimeUnit.SECONDS)) {
                taskExecutor.shutdownNow();
            }
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            taskExecutor.shutdownNow();
        }
        
        log.info("Process management service shutdown completed");
    }
} 